"""
src/basebody.py
    Create a Base Body object class derived from the ode Body.  This
    class provides all the hooks necessary to have a fully featured
    Body without a lot of work by the user.

Written By:
    James Thomas
    Email: jim@houseoftechnology.org
    Web: http://mission-cognition.houseoftechnology.org/

Copyright 2009-2012

This file is part of the PandaODElib distribution.

This library is free software; you can redistribute it and/or
modify it under the terms of the GNU Lesser General Public License
as published by the Free Software Foundation; either version 3 of
the License, or (at your option) any later version. The text of the
GNU Lesser General Public License is included with this library in
the file named LICENSE.

This library is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
"""
from pandac.PandaModules import Quat, Vec3, Vec3D
from pandac.PandaModules import OdeBody, OdeFixedJoint

def QuatInverse(Q):
    """Panda3d inverse does not conform to any definition I've seen.  Use this instead."""
    return Q.conjugate() / Q.lengthSquared()

def QuatMultiply(Q1, Q2):
    """Panda3d multiplication order seems backwards from all the definitions I've seen. Use this instead"""
    return Q2 * Q1

def QuatFromQuaternionPair(q1, q2): 
    """Creates a quaternion q which will rotate q1 to q2 such that RotateQuat(q, q1) -> q2"""
    return QuatMultiply(q2, QuatInverse(q1))

class BaseBody(OdeBody):
    # Provide new so we can have subclasses that add parameters
    def __new__(cls, *args, **kwargs):
        # Initialize the Body class
        return OdeBody.__new__(cls, args[0]) # First arg is World

    def __init__(self, World):
        OdeBody.__init__(self, World)
        
        # Remember the world so we don't need it passed everywhere
        self._world = World

        # Initialize the connections as none defined.
        # The connections can be made later by calling the appropriate functions.
        self._connections = {} # Use a dictionary so we don't have to use an index for ID
        self._connectionLocations = {}

    def _setConnectionJointInstance(self, ConnectionID, JointInstance):
        self._connections[ConnectionID] = JointInstance

    def RotateOrientation(self, AngleRad, Axis):
        """Rotate the body orientation by the angle about the specified axis (must be unit vector).
        This rotation does not result any any translation of the body location.
        The axis is specified in the world coordinate system and the angle is in radians."""
        axis = Vec3(Axis)
        axis.normalize()
        qr = Quat()
        qr.setFromAxisAngleRad(AngleRad, axis)
        self.setQuaternion(QuatMultiply(qr, Quat(self.getQuaternion())))

    def RotateAboutPoint(self, Point, Quaternion):
        """Rotates the body about the specified point by the rotation quaternion (normalized).
        the location of the body is moved as a result of the rotation.
        The point and quaternion are specified in the world coordinate system."""

        point = Vec3(Point)
        qr = Quat(Quaternion)
        
        # Get the body position relative to the Point
        relPos = self.getPosition() - point

        # Rotate the cg-relative position by the quaternion
        newRelPos = qr.xform(relPos)
        
        # Move the body to the new absolute location
        self.setPosition(newRelPos + Vec3(Point))
   
        # A rotation is equivalent to a quaternion multiply: qNewOrientation = qRotate * qOldOrientation
        qNew = QuatMultiply(qr, Quat(self.getQuaternion()))
        
        # Set the new orientation
        self.setQuaternion(qNew)

    def MoveByDelta(self, Delta):
        """Translate the body position by Delta 3-tuple vector."""
        self.setPosition(self.getPosition() + Vec3(Delta))

    def MoveToPlaceConnectionAt(self, ConnectionID, Position):
        """Move the body such that the specified connection point is located at the specified
        3-tuple position vector.  The position is specified in world coordinates."""
        assert ConnectionID in self._connectionLocations

        # Get the location of the specified connection point.
        location = self.GetConnectionLocation(ConnectionID)

        # Calculate the delta between the new position and the connection point
        delta = Vec3(Position) - Vec3(location)

        # Move the body
        self.MoveByDelta(delta)

    def OrientConnectionAt(self, ConnectionID, Orientation):
        """Orient the body such that the specified connection orientation is aligned with the
        specified orientation quaternion.  The quaternion is specified in world coordinates."""
        assert ConnectionID in self._connectionLocations

        # Get the axis of the specified connection point (in world coordinates).
        connQ = self.GetConnectionOrientation(ConnectionID)

        qr = QuatFromQuaternionPair(connQ, Orientation)

        # Get the location of the specified connection point.
        location = self.GetConnectionLocation(ConnectionID)

        # Rotate the body about the connection point
        self.RotateAboutPoint(location, qr)

    def DefineConnectionPoint(self, ConnectionID, Position, Orientation=Quat()):
        """Define a connection point on the body for use in placing joints.  The position
        is a 3-tuple vector in world coordinates.  The connection point can have an optional
        orientation (specified in world coordinates) which may be used by a joint that requires
        a particular orientation axis/angle."""
        # For now we cannot deal with 'moving' a location
        assert ConnectionID not in self._connectionLocations

        # Convert the position and orientation to local coordinates before storing
        localPos = self.getPosRelPoint(Position)

        localQ = QuatFromQuaternionPair(Quat(self.getQuaternion()), Orientation)

        self._connectionLocations[ConnectionID] = (localPos, localQ)

    def Connect(self, JointInstance, ConnectionID, Other=None, OtherConnectionID=''):
        """Make a connection using the specified joint instance between the specified connection
        on this body and the specified connection on the other body/assembly. If Other is not 
        specified (or None) then the connection is made to the environment."""
        assert ConnectionID in self._connectionLocations
        assert ConnectionID not in self._connections

        if Other is not None:
            # Orient Other Body such that the connections between the two bodies line up
            pos = self.GetConnectionLocation(ConnectionID)
            Other.MoveToPlaceConnectionAt(OtherConnectionID, pos)

            q = self.GetConnectionOrientation(ConnectionID)
            Other.OrientConnectionAt(OtherConnectionID, q)
            
            if hasattr(Other, 'GetBodyFromConnectionID'):
                otherBody = Other.GetBodyFromConnectionID(OtherConnectionID)
            elif isinstance(Other, OdeBody):
                otherBody = Other
            else:
                raise TypeError, 'Unexpected type for Other parameter: ' + type(Other)
        else:
            otherBody = None
            
        # Make the connection.
        # If OtherBody is None then it attaches to the environment
        JointInstance.attach(self, otherBody)

        # Get the information for this joint location
        location = self.GetConnectionLocation(ConnectionID)
        orientation = self.GetConnectionOrientation(ConnectionID)

        # Set the location of the joint to the prepared location
        try:
            JointInstance.setAnchor(location)
        except AttributeError:
            pass

        try:
            # Compute the joint axis based on the connection orientation
            axis = orientation.xform((0,0,1))
            JointInstance.setAxis(axis)
        except AttributeError:
            # Not all joints have a setAxis() method
            pass

        # Put the joint object into this body and the other body/assembly connection lists
        self._setConnectionJointInstance(ConnectionID, JointInstance)
        if Other is not None:
            Other._setConnectionJointInstance(OtherConnectionID, JointInstance)

    def IsConnected(self, ConnectionID):
        """Returns True if the specified connectionID has a joint connected."""
        return ConnectionID in self._connections

    def GetConnectionJoint(self, ConnectionID):
        """Gets the joint instance for the specified connection.  Use IsConnected() to 
        make sure the connection has a joint before calling this function."""
        assert ConnectionID in self._connections

        if ConnectionID in self._connections:
            # Get the joint instance for this connection
            return self._connections[ConnectionID]
        else:
            return None

    def GetConnectionIDFromJoint(self, JointInstance):
        """Call this function when you have the joint instance but don't know where it is
        connected on the body.  returns None if joint is not connected on this body.""" 
        keys = [key for key, value in self._connections.items() if value == JointInstance]
        if keys:
            connectionID = keys[0]
        else:
            connectionID = None
        return connectionID

    def Disconnect(self, ConnectionID):
        """Disconnect the joint at this connection and any connection on another body.
        The joint instance will be lost if it is not held elsewhere.  Use IsConnected() to 
        make sure the connection has a joint before calling this function."""
        # Get the joint instance
        jointInstance = self.GetConnectionJoint(ConnectionID)

        # Figure out the other body
        otherBody = jointInstance.getBody(0)
        if otherBody == self:
            otherBody = jointInstance.getBody(1)

        # Figure out the connection id on the other body
        otherConnectionID = otherBody.GetConnectionIDFromJoint(jointInstance)

        if otherConnectionID is not None:
            # Disconnect the other body
            otherBody._Connections.pop(otherConnectionID)

        # Disconnect the joint
        self._connections[ConnectionID].attach(None, None)

        # Now disconnect ourselves
        self._connections.pop(ConnectionID)

    def SetFixed(self):
        """Fix this body to the ode world using a FixedJoint."""
        self._fixedJoint = OdeFixedJoint(self._world)
        self._fixedJoint.attach(self, None) # None sets it to the environment
        self._fixedJoint.set()

    def ReleaseFixed(self):
        """Release a fixed joint from SetFixed() if one was specified for this body."""
        try:
            self._fixedJoint.detach()
            self._fixedJoint.destroy()
            del self._fixedJoint
        except AttributeError:
            pass

    def GetConnectionLocation(self, ConnectionID):
        """Get the connection point location in world coordinates."""
        assert ConnectionID in self._connectionLocations
        return self.getRelPointPos(self._connectionLocations[ConnectionID][0])

    def GetConnectionLocationIDs(self):
        return self._connectionLocations.keys()
    
    def GetConnectionIDs(self):
        """Get a list of the IDs of any specified connections."""
        return self._connections.keys()

    def GetConnectionOrientation(self, ConnectionID):
        """Get the orientation of a connection in world coordinates."""
        assert ConnectionID in self._connectionLocations

        localQ = self._connectionLocations[ConnectionID][1]
        
        # Convert quaternion to global coordinates
        return QuatMultiply(Quat(self.getQuaternion()), localQ)

    def GetAngularMomentum(self):
        """Returns the vector L = Iw"""
        Ix, Iy, Iz = self.getMass().getInertialTensor() # Body local coordinates    
        wx, wy, wz = self.vectorFromWorld(self.getAngularVel()) # Convert to local
        Lx = Ix[0] * wx + Ix[1] * wy + Ix[2] * wz
        Ly = Iy[0] * wx + Iy[1] * wy + Iy[2] * wz
        Lz = Iz[0] * wx + Iz[1] * wy + Iz[2] * wz
        
        return (Lx, Ly, Lz)
    
    def UpdateSettings(self):
        """This function must be called before each time step is executed.  It will
        call the UpdateSettings() function for any connected joints that have it."""
        # Update the joints
        for jointInstance in self._connections.itervalues():
            # This may seem convoluted but I do it this way on purpose.
            # When I originally called the function in the try: the exception
            # was catching errors in the implementation of the update function.
            # This is made it a pain to debug so now I do it in two passes.
            updateFunc = None
            try:
                updateFunc = jointInstance.UpdateSettings
            except AttributeError:
                # The ode joints don't have an UpdateSettings function
                pass

            # If updateFunc is still None then obviously it didn't exist
            if updateFunc is not None:
                updateFunc()
